from django.shortcuts import render
from requests.api import request
from rest_framework.response import Response
from rest_framework.viewsets import GenericViewSet
from rest_framework.decorators import action
from rest_framework.permissions import AllowAny
from rest_framework.exceptions import APIException
from rest_framework import status
# Create your views here.
from pcap_jobs.models import NodeInfo
import rsa
import base64
class PcapRegister(GenericViewSet):
    @action(methods=["post"],detail=False)
    def register(self,request):
        """
        节点注册接口
        """
        req = request.data 
        _ip = req.get('ip')
        if _ip is None:
            raise APIException("请输入ip地址",status.HTTP_200_OK)

        node = NodeInfo.objects.filter(host_id = _ip)
        if len(node) != 0:
            raise APIException("此IP已注册,请勿再次注册",status.HTTP_200_OK)

        # 生成密钥
        (pubkey, privkey) = rsa.newkeys(1024)
        publickey = pubkey.save_pkcs1().decode()
        privatekey = privkey.save_pkcs1().decode()
        try:
            nodeInfo = NodeInfo(host_ip=_ip,public_key=publickey,private_key=privatekey)   
            nodeInfo.save()
        except:
            raise APIException("保存失败",status.HTTP_200_OK)
        return Response({"pubkey":publickey})

    @action(methods=["post"],detail=False)
    def verifySign(self,request):
        """
        用于验证节点发送的挑战或者反馈签名
        """
        req = request.data 
        _ip = req.get("ip")
        if _ip is None:
            raise APIException("请输入ip",status.HTTP_200_OK)
        _pcap_path = req.get("pcap")
        if _pcap_path is None:
            raise APIException("请输入pcap地址",status.HTTP_200_OK)
        _sign = req.get("sign")
        if _sign is None:
            raise APIException("请输入签名",status.HTTP_200_OK)
        pubkey_node = NodeInfo.objects.filter(host_id=_ip).first()
        if pubkey_node is None:
            raise APIException("此节点未注册",status.HTTP_200_OK)
        pubkey = pubkey_node.public_key
        with open(_pcap_path, 'rb') as p:
            pcap = p.read()
            pcap = base64.b64encode(pcap)
            # 公钥验证
            if rsa.verify(pcap, _sign, pubkey):
                return Response("OK")
            else:
                return Response("faile")
                